from ryu.base import app_manager
from ryu.controller import ofp_event
from ryu.controller.handler import MAIN_DISPATCHER, set_ev_cls
from ryu.ofproto import ofproto_v1_0
from ryu.lib.packet import packet, ethernet, ipv4

class CustomSwitch(app_manager.RyuApp):
    OFP_VERSIONS = [ofproto_v1_0.OFP_VERSION]

    def __init__(self, *args, **kwargs):
        super(CustomSwitch, self).__init__(*args, **kwargs)
    
    @set_ev_cls(ofp_event.EventOFPPacketIn, MAIN_DISPATCHER)
    def _packet_in_handler(self, ev):
        msg = ev.msg
        datapath = msg.datapath
        ofproto = datapath.ofproto
        parser = datapath.ofproto_parser
        dpid = datapath.id
        
        self.logger.info("Packet in DPID: %s", dpid)  # 打印 DPID
        pkt = packet.Packet(msg.data)
        eth = pkt.get_protocol(ethernet.ethernet)
        ip_pkt = pkt.get_protocol(ipv4.ipv4)

        if ip_pkt is None:
            return

        src_ip = ip_pkt.src
        dst_ip = ip_pkt.dst
        
        self.logger.info("Received packet: SRC IP: %s, DST IP: %s", src_ip, dst_ip)

        # 根据源 IP 和目的 IP 来决定流表规则
        if dpid == 1:  # Switch 1
            if src_ip == '192.168.0.101' and dst_ip == '192.168.0.104':  # h1 -> h4
                out_port = 4
            elif src_ip == '192.168.0.104' and dst_ip == '192.168.0.101':  # h4 -> h1
                out_port = 1
            elif src_ip == '192.168.0.102' and dst_ip == '192.168.0.105':  # h2 -> h5
                out_port = 4
            elif src_ip == '192.168.0.105' and dst_ip == '192.168.0.102':  # h5 -> h2
                out_port = 2
            elif src_ip == '192.168.0.103' and dst_ip == '192.168.0.106':  # h3 -> h6
                out_port = 4
            elif src_ip == '192.168.0.106' and dst_ip == '192.168.0.103':  # h6 -> h3
                out_port = 3
            else:
                return
        elif dpid == 2:  # Switch 2
            if src_ip == '192.168.0.101' and dst_ip == '192.168.0.104':  # h1 -> h4
                out_port = 1
            elif src_ip == '192.168.0.104' and dst_ip == '192.168.0.101':  # h4 -> h1
                out_port = 4
            elif src_ip == '192.168.0.102' and dst_ip == '192.168.0.105':  # h2 -> h5
                out_port = 2
            elif src_ip == '192.168.0.105' and dst_ip == '192.168.0.102':  # h5 -> h2
                out_port = 4
            elif src_ip == '192.168.0.103' and dst_ip == '192.168.0.106':  # h3 -> h6
                out_port = 3
            elif src_ip == '192.168.0.106' and dst_ip == '192.168.0.103':  # h6 -> h3
                out_port = 4
            else:
                return
        else:
            return
           
        self.logger.info("Installing flow: DPID %s, SRC IP %s -> DST IP %s, OUT PORT %s", dpid, src_ip, dst_ip, out_port)

        actions = [parser.OFPActionOutput(out_port)]

        # 安装流表规则
        match = parser.OFPMatch(dl_type=eth.ethertype, nw_src=src_ip, nw_dst=dst_ip)
        mod = parser.OFPFlowMod(
            datapath=datapath, priority=1, match=match, actions=actions)
        datapath.send_msg(mod)

        # 立即转发数据包
        out = parser.OFPPacketOut(
            datapath=datapath, buffer_id=msg.buffer_id,
            in_port=msg.in_port, actions=actions, data=msg.data)
        datapath.send_msg(out)

